#include "proto.h"
#include "pbc.h"
#include "parser.h"
#include "number.h"
#include "str.h"
#include "hashmap.h"
#include "array.h"
#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <string.h>
#include <strings.h>

#define WMESSAGE_SIZE 128

#define ARRAY_SET_UZZ_TEMPLATE(m, ptr, type_value, _type, decode_func, zigzag_func)\
do{\
    int ret, len = 0;\
    ret = pbc_decode32(ptr, &len);\
    ptr += ret;\
    int size = m->size++;\
    struct simple_value *sv = &m->slot[size];\
    sv->slot = malloc(sizeof(value) * len);\
    strcpy(sv->name, f->name);\
    int i, jump = 0;\
    for (i = 0; i < len; ++i)\
    {\
        _type v;\
        value *vs = &(sv->slot[i]);\
        _type n = decode_func(ptr, &v);\
        vs->v.type_value = zigzag_func(v);\
        ptr += n;\
        jump += n;\
        if (jump == len)\
            break;\
    }\
    sv->size = ++i;\
    return ret + len;\
} while (0);


#define ARRAY_SET_NZZ_TEMPLATE(m, ptr, type_value, _type, decode_func)\
do{\
    int ret, len = 0;\
    ret = pbc_decode32(ptr, &len);\
    ptr += ret;\
    int size = m->size++;\
    struct simple_value *sv = &m->slot[size];\
    sv->slot = malloc(sizeof(value) * len);\
    strcpy(sv->name, f->name);\
    int i, jump = 0;\
    for (i = 0; i < len; ++i)\
    {\
        value *vs = &(sv->slot[i]);\
        _type n = decode_func(ptr, &(vs->v.type_value));\
        ptr += n;\
        jump += n;\
        if (jump == len)\
            break;\
    }\
    sv->size = ++i;\
    return ret + len;\
} while (0);


#define GET_VALUE_TYPE(rmsg, fname, type, idx)\
do{\
    value *vs = getvalue(rmsg, fname, idx);\
    if (vs)\
    {\
        return vs->v.type;\
    }\
} while (0);


typedef union 
{
    int i32;
    int64_t i64;
    float f;
    double d;
    char *str;
    struct pbc_rmessage *rm;
}base_type;

typedef struct _value
{
    base_type v;
    int isalloc;
    char name[NAME_SIZE];
}value;

struct simple_value
{
    int size;
    value *slot;
    char name[NAME_SIZE];
};

struct pbc_wmessage
{
    struct pbc_env *env;
    struct pbc_wmessage *ref;    //自身的引用
    struct pbc_wmessage *tail;
    struct pbc_field *field;
    struct repeat_array *array;

    uint8_t *buffer;
    uint8_t *ptr;
    int cap;
};

struct pbc_rmessage
{
    struct simple_value sv;
    struct pbc_env *env;
    struct simple_value *slot;
    int size;
};

static void pbc_rmessage_buffer(struct pbc_rmessage *m, struct chumk *ck);
static struct pbc_rmessage *rmessage_decode(struct pbc_env *env, const char *msgname, struct chumk *ck);



struct pbc_wmessage *pbc_wmessage_create(struct pbc *p, const char *msgname)
{
    struct pbc_env *env = pbc_env_query(p, msgname);
    if (env == NULL)
    {
        fprintf(stderr, "can't find message(%s).\n", msgname);
        exit(1);
    }

    if (env->field_sz == 0)
    {
        fprintf(stderr, "message(%s) field size is zero.\n", msgname);
        exit(1);   
    }

    int sz = sizeof(struct pbc_wmessage);
    struct pbc_wmessage *wmsg = malloc(sz);
    wmsg->ref = NULL;
    wmsg->tail = NULL;
    wmsg->env = env;
    wmsg->cap = WMESSAGE_SIZE;
    wmsg->buffer = malloc(wmsg->cap);
    wmsg->ptr = wmsg->buffer;
    wmsg->array = malloc(repeat_array_block());
    repeat_array_init(wmsg->array);

    return wmsg;
}

void pbc_wmessage_free(struct pbc_wmessage *m)
{
    if (m->ref == NULL)
    {
        repeat_array_free(m->array);
        free(m->array);
        free(m->buffer);
        free(m);
        return;
    }
    struct pbc_wmessage *next = NULL;
    while (m)
    {
        next = m->ref;
        repeat_array_free(m->array);
        free(m->array);
        free(m->buffer);
        free(m);
        m = next;
    }
}

struct pbc_wmessage *expand_wmessage(struct pbc_wmessage *wmsg)
{
    int sz = wmsg->ptr - wmsg->buffer;
    int cap = wmsg->cap;
    if (sz > cap * 3/4)
    {
        uint8_t *buffer = malloc(cap * 2);
        memcpy(buffer, wmsg->buffer, sz);
        free(wmsg->buffer);
        wmsg->buffer = buffer;
        wmsg->ptr = wmsg->buffer + sz;
        wmsg->cap = cap * 2;
    }

    return wmsg;
}

void pbc_wmessage_buffer(struct pbc_wmessage *m, struct chumk *ck)
{
    struct chumk cck = {m->ptr - m->buffer, m->buffer};
    struct chumk *tmp = combine_wmessage_array(m->array, &cck, &m->cap);
    m->buffer = tmp->buffer;
    m->ptr = m->buffer + tmp->len;

    if(m->ref == NULL)
    {
        ck->buffer = m->buffer;
        ck->len = m->ptr - m->buffer;
        return;
    }
    
    int len = m->ptr - m->buffer;
    int sz = len < 128 ? 128 : len * 2;
    int cap = sz;

    uint8_t *buffer = malloc(cap);
    uint8_t *ptr = buffer;

    memcpy(ptr, m->buffer, len);
    ptr += len;

    struct pbc_wmessage *next = m->ref;
    while(next)
    {
        int n = next->ptr - next->buffer;

        if (n > sz - 5)
        {
            cap *= 2;
            len = ptr - buffer;
            uint8_t *nbuffer = malloc(cap);
            memcpy(nbuffer, buffer, len);

            free(buffer);

            buffer = nbuffer;
            ptr = buffer + len;

            assert(len == sz);
            sz = cap - sz;
        }

        int tag = next->field->field_number << 3 | WIRETYPE_LENGTH;
        len = pbc_encode32(tag, ptr);
        sz -= len;
        ptr += len;

        len = pbc_encode32(n, ptr);
        sz -= len;
        ptr += len;

        memcpy(ptr, next->buffer, n);
        sz -= n;
        ptr += n;
        
        next = next->ref;
    }

    free(m->buffer);

    m->buffer = buffer;
    m->ptr = ptr;

    ck->buffer = m->buffer;
    ck->len = m->ptr - m->buffer;
}

struct pbc_wmessage *pbc_wmessage_message(struct pbc_wmessage *m, const char *msgname)
{
    struct pbc_field *f = pbc_field_query(m->env, msgname);
    assert(f);

    struct pbc_env *env = f->m;
    if (env == NULL)
    {
        fprintf(stderr, "can't find message(%s).\n", msgname);
        exit(1);
    }

    if (env->field_sz == 0)
    {
        fprintf(stderr, "message(%s) field size is zero.\n", msgname);
        exit(1);
    }

    int sz = sizeof(struct pbc_wmessage);
    struct pbc_wmessage *wmsg = malloc(sz);

    if (m->ref == NULL)
    {
        m->ref = wmsg;
    }
    else
    {
        m->tail->ref = wmsg;
    }

    m->tail = wmsg;
    wmsg->env = env;
    wmsg->cap = WMESSAGE_SIZE;
    wmsg->buffer = malloc(wmsg->cap);
    wmsg->ptr = wmsg->buffer;
    wmsg->ref = NULL;
    wmsg->field = f;
    wmsg->array = malloc(repeat_array_block());
    repeat_array_init(wmsg->array);
    
    return wmsg;
}

int pbc_encode_int32(struct pbc_wmessage *wmsg, const char *fieldname, int value)
{
    struct pbc_field *f = pbc_field_query(wmsg->env, fieldname);
    if (f == NULL)
    {
        fprintf(stderr, "error: message %s can't find field: %s\n", wmsg->env->name, fieldname);
        return -1;
    }

    if (f->wire_type != PBC_INT32 && f->wire_type != PBC_SINT32)
    {
        fprintf(stderr, "encode %s type error: can't switch %s to int32 or sint32\n", 
                fieldname, wire_to_string(f->wire_type));
        return -1;
    }

    wmsg = expand_wmessage(wmsg);
    uint8_t *startoff = wmsg->ptr;
    
    if (f->wire_type == PBC_SINT32)
    {
        value = int_to_zigzag(value);
    }

    if (f->dec_type == 0 || f->dec_type == DECORATE_OPTIONAL)
    {
        int tag = f->field_number << 3 | WIRETYPE_VARINT;
        int sa = pbc_encode32(tag, wmsg->ptr);
        wmsg->ptr += sa;
        wmsg->ptr += pbc_encode32(value, wmsg->ptr);
    }
    else if (f->dec_type == DECORATE_REPEATED)
    {
        int tag = f->field_number << 3 | WIRETYPE_LENGTH;
        repeat_array_pushi32(wmsg->array, fieldname, tag, value);
    }

    return wmsg->ptr - startoff;
}

int pbc_encode_int64(struct pbc_wmessage *wmsg, const char *fieldname, int64_t value)
{
    struct pbc_field *f = pbc_field_query(wmsg->env, fieldname);
    assert(f);

    if (f->wire_type != PBC_INT64 && f->wire_type != PBC_SINT64)
    {
        fprintf(stderr, "encode %s type error: can't switch %s to int64 or sint64\n", 
                fieldname, wire_to_string(f->wire_type));
        return -1;
    }

    wmsg = expand_wmessage(wmsg);
    uint8_t *startoff = wmsg->ptr;
    
    if (f->wire_type == PBC_SINT64)
    {
        value = int64_to_zigzag(value);
    }
    
    if (f->dec_type == 0 || f->dec_type == DECORATE_OPTIONAL)
    {
        int tag = f->field_number << 3 | WIRETYPE_VARINT;
        wmsg->ptr += pbc_encode32(tag, wmsg->ptr);
        wmsg->ptr += pbc_encode64(value, wmsg->ptr);
    }
    else if (f->dec_type == DECORATE_REPEATED)
    {
        int tag = f->field_number << 3 | WIRETYPE_LENGTH;
        repeat_array_pushi64(wmsg->array, fieldname, tag, value);
    }

    return wmsg->ptr - startoff;
}

int pbc_encode_float(struct pbc_wmessage *wmsg, const char *fieldname, float value)
{
    struct pbc_field *f = pbc_field_query(wmsg->env, fieldname);
    assert(f);

    if (f->wire_type != PBC_FLOAT)
    {
        fprintf(stderr, "encode %s type error: can't switch %s to float\n", 
                fieldname, wire_to_string(f->wire_type));
        return -1;
    }

    wmsg = expand_wmessage(wmsg);
    uint8_t *startoff = wmsg->ptr;
    
    if (f->dec_type == 0 || f->dec_type == DECORATE_OPTIONAL)
    {
        int tag = f->field_number << 3 | WIRETYPE_FIXED32;
        wmsg->ptr += pbc_encode32(tag, wmsg->ptr);
        wmsg->ptr += pbc_encodef(value, wmsg->ptr);
    }
    else if (f->dec_type == DECORATE_REPEATED)
    {
        // int tag = f->field_number << 3 | WIRETYPE_LENGTH;
        // repeat_array_pushi64(wmsg->array, fieldname, tag, value);
    }

    return wmsg->ptr - startoff;
}

int pbc_encode_string(struct pbc_wmessage *wmsg, const char *fieldname, const char *value)
{
    struct pbc_field *f = pbc_field_query(wmsg->env, fieldname);
    assert(f);
    
    wmsg = expand_wmessage(wmsg);

    uint8_t *startoff = wmsg->ptr;
    if (f->dec_type == 0 || f->dec_type == DECORATE_OPTIONAL)
    {
        if (f->wire_type == PBC_STRING)
        {
            int tag = f->field_number << 3 | WIRETYPE_LENGTH;
            int len = strlen(value);
            wmsg->ptr += pbc_encode32(tag, wmsg->ptr);
            wmsg->ptr += pbc_encode32(len, wmsg->ptr);

            memcpy(wmsg->ptr, value, len);
            wmsg->ptr += len;
        }
    }
    else if (f->dec_type == DECORATE_REPEATED)
    {
        int tag = f->field_number << 3 | WIRETYPE_LENGTH;
        repeat_array_pushstr(wmsg->array, fieldname, tag, value);
    }

    return wmsg->ptr - startoff;
}


static value *getvalue(struct pbc_rmessage *m, const char *name, int idx)
{
    int i;
    struct pbc_field *f = pbc_field_query(m->env, name);
    if (f == NULL)
    {
        printf("can't find the field(%s) of message(%s)\n", name, m->env->name);
        return NULL;
    }

    if (f->dec_type == DECORATE_REPEATED)
    {
        for (i = 0; i < m->size; ++i)
        {
            struct simple_value *sv = &m->slot[i];
            if (strcmp(sv->name, name) == 0)
            {
                if (idx >= 0 && idx < sv->size)
                    return &sv->slot[idx];
            }
        }
    }
    else
    {
        struct simple_value *sv = &(m->sv);
        for (i = 0; i < sv->size; ++i)
        {
            if (strcmp(sv->slot[i].name, name) == 0)
                return &sv->slot[i];
        }
    }

    return NULL;
}

static int decode_type_varint(int ftype, value *vs, char *name, uint8_t *ptr)
{
    int ret = 0;

    switch(ftype)
    {
        case PBC_INT32:
            ret = pbc_decode32(ptr, &(vs->v.i32));
            break;
        case PBC_SINT32:
            ret = pbc_decode32(ptr, &(vs->v.i32));
            vs->v.i32 = zigzag_to_int(vs->v.i32);
            break;
        case PBC_INT64:
            ret = pbc_decode64(ptr, &(vs->v.i64));
            break;
        case PBC_SINT64:
            ret = pbc_decode64(ptr, &(vs->v.i64));
            vs->v.i64 = zigzag_to_int64(vs->v.i64);
            break;
        default:
            fprintf(stderr, "decode_type_int error\n");
            break;
    }

    sprintf(vs->name, "%s", name);
    return ret;
}


static int decode_type_length(int ftype, value *vs, struct pbc_field *f, uint8_t *ptr, struct pbc_rmessage *m, int *simple)
{
    int ret = 0;

    switch(ftype)
    {
        case PBC_STRING:
        {
            int len = 0;
            if (f->dec_type != DECORATE_REPEATED)
            {
                *simple = 1;
                ret = pbc_decode32(ptr, &len);
                ptr += ret;
                vs->v.str = strdump((const char*)ptr, len);
                vs->isalloc = 1;
                ret += len;
            }
            else if (f->dec_type == DECORATE_REPEATED)
            {

            }
            break;
        }
        case PBC_MESSAGE:
        {
            *simple = 1;

            int len = 0;
            ret = pbc_decode32(ptr, &len);
            ptr += ret;

            struct chumk ck = {len, ptr};
            struct pbc_rmessage *rm = rmessage_decode(f->m, f->name, &ck);

            vs->v.rm = rm;
            vs->isalloc = 2;
            ret += len;
            break;
        }
        case PBC_SINT32:
        {
            ARRAY_SET_UZZ_TEMPLATE(m, ptr, i32, int, pbc_decode32, zigzag_to_int);
        }
        case PBC_SINT64:
        {
            ARRAY_SET_UZZ_TEMPLATE(m, ptr, i64, int64_t, pbc_decode64, zigzag_to_int64);
        }
        case PBC_INT32:
        {
            ARRAY_SET_NZZ_TEMPLATE(m, ptr, i32, int, pbc_decode32);
        }
        case PBC_INT64:
        {
            ARRAY_SET_NZZ_TEMPLATE(m, ptr, i64, int64_t, pbc_decode64);
        }
        default:
        {
            fprintf(stderr, "decode_type_length error, ftype: %d\n", ftype);
            exit(1);
        }
    }
    sprintf(vs->name, "%s", f->name);
    return ret;
}

static int decode_type_fix(int ftype, value *vs, char *name, uint8_t *ptr, struct pbc_rmessage *m) 
{
    int ret = 0;
    switch(ftype)
    {
        case PBC_FLOAT:
        {
            union {
                float f;
                uint32_t e;
            } u;

            u.e = pbc_decodeu32(ptr);
            vs->v.f = u.f;
            ret = 4;
            break;
        }
    }

    sprintf(vs->name, "%s", name);
    return ret;
}

static int decode_one_field(int tag, struct pbc_rmessage *m, uint8_t *ptr)
{
    int fnum = tag >> 3;
    int wire = tag & 0x07;
    int simple = 0;
    uint8_t *off = ptr;
    
    value v;
    bzero(&v, sizeof(v));

    struct pbc_field *f = pbc_field_query_fn(m->env, fnum);
    int ftype = f->wire_type;

    if (wire == WIRETYPE_VARINT)
    {
        simple = 1;
        ptr += decode_type_varint(ftype, &v, f->name, ptr);
    }
    else if (wire == WIRETYPE_LENGTH)
    {
        ptr += decode_type_length(ftype, &v, f, ptr, m, &simple);   
    }
    else if (wire == WIRETYPE_FIXED32)
    {
        simple = 1;
        ptr += decode_type_fix(ftype, &v, f->name, ptr, m);  
    }

    // 简单的数据类型，才保存到 m->sv.slot　中
    if (f->dec_type != DECORATE_REPEATED && simple)
    {
        struct pbc_env *env = m->env;
        int cap = env->field_sz - env->array_sz;

        if (m->sv.size >= cap * 8)
        {
            fprintf(stderr, "error size:%d, cap:%d, name:%s\n", m->sv.size, cap, env->name);
            exit(1);
        }
        m->sv.slot[m->sv.size++] = v;
    }

    return ptr - off;
}

static void pbc_rmessage_buffer(struct pbc_rmessage *m, struct chumk *ck)
{
    uint8_t *ptr = ck->buffer;
    int len = ck->len;

    while (len > 0)
    {
        int tag = 0;
        int n = pbc_decode32(ptr, &tag);
        len -= n;
        ptr += n;

        n = decode_one_field(tag, m, ptr);
        len -= n;
        ptr += n;
    }
}

static struct pbc_rmessage *rmessage_decode(struct pbc_env *env, const char *msgname, struct chumk *ck)
{
    if (env == NULL)
    {
        fprintf(stderr, "2can't find message(%s).\n", msgname);
        exit(1);
    }

    if (env->field_sz == 0)
    {
        fprintf(stderr, "2message(%s) field size is zero.\n", msgname);
        exit(1);   
    }

    int ssz = env->field_sz - env->array_sz;
    struct pbc_rmessage *m = malloc(sizeof(*m));
    m->env = env;
    
    m->sv.slot = NULL;
    m->sv.size = 0;
    
    m->slot = NULL;
    m->size = 0;
    
    if (ssz > 0)
    {
        m->sv.slot = malloc(sizeof(value) * ssz * 8);
    }

    if (env->array_sz > 0)
    {
        m->slot = malloc(sizeof(struct simple_value) * env->array_sz);
    }

    pbc_rmessage_buffer(m, ck);
    return m;
}

struct pbc_rmessage *pbc_rmessage_create(struct pbc *p, const char *msgname, struct chumk *ck)
{
    struct pbc_env *env = pbc_env_query(p, msgname);
    struct pbc_rmessage *m = rmessage_decode(env, msgname, ck);
    return m;
}

void pbc_rmessage_free(struct pbc_rmessage *m)
{
    int i, sz = m->env->field_sz;
    value *vs = m->sv.slot;
    for (i = 0; i < sz; ++i)
    {
        if (vs[i].isalloc == 1)
            free(vs[i].v.str);
        else if (vs[i].isalloc == 2)
            pbc_rmessage_free(vs[i].v.rm);
    }

    if (m->size > 0)
    {
        for (i = 0; i < m->size; ++i)
        {
            free(m->slot[i].slot);
        }
    }

    free(m->slot);
    free(m->sv.slot);
    free(m);
}

struct pbc_rmessage *pbc_rmessage_message(struct pbc_rmessage *m, const char *msgname, int idx)
{
    int find = -1;
    int i, sz = m->sv.size;
    value *vs = m->sv.slot;
    for ( i = 0; i < sz; ++i)
    {
        if (strcmp(vs[i].name, msgname) == 0)
        {
            find++;
            if (find == idx)
            {
                return vs[i].v.rm;
            }
        }
    }

    return NULL;
}

int pbc_decode_int32(struct pbc_rmessage *rmsg, const char *fieldname, int idx, int def)
{
    GET_VALUE_TYPE(rmsg, fieldname, i32, idx);
    return def;
}

int64_t pbc_decode_int64(struct pbc_rmessage *rmsg, const char *fieldname, int idx, int64_t def)
{
    GET_VALUE_TYPE(rmsg, fieldname, i64, idx);
    return def;
}

float pbc_decode_float(struct pbc_rmessage *rmsg, const char *fieldname, int idx, float def)
{
    GET_VALUE_TYPE(rmsg, fieldname, f, idx);
    return def;
}

const char *pbc_decode_string(struct pbc_rmessage *rmsg, const char *fieldname, int idx, const char *def)
{
    GET_VALUE_TYPE(rmsg, fieldname, str, idx);
    return def;
}

int pbc_rmessage_size(struct pbc_rmessage *m, const char *fieldname)
{
    int i;
    struct pbc_field *f = pbc_field_query(m->env, fieldname);
    if (f == NULL)
    {
        return 0;
    }
    else if (f->dec_type == DECORATE_REPEATED)
    {
        for (i = 0; i < m->size; ++i)
        {
            struct simple_value *sv = &m->slot[i];
            if (strcmp(sv->name, fieldname) == 0)
                return sv->size;
        }
    }

    int counter = 0;
    int sz = m->sv.size;
    value *vs = m->sv.slot;
    for ( i = 0; i < sz; ++i)
    {
        if (strcmp(vs[i].name, fieldname) == 0)
            counter++;
    }

    if (counter)
    {
        return counter;
    }

    return 1;
}


